package com.bigdata.kafka.utils;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.DefaultParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.MissingArgumentException;
import org.apache.commons.cli.Options;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.DateUtils;
import org.apache.commons.lang3.time.FastDateFormat;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.sql.Timestamp;
import java.text.SimpleDateFormat;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;


/**
 * kafka 消费工具类
 * -bs hadoop:9092 -topic test -time -2 -i 1h -pt -d
 * <p>
 * -bs docker5:7051 -topic test1 -time -2 -i 1h -pt -d
 *
 * @author xin.ding
 * @date 2020/8/21 11:24
 */
public class TimestampConsumer {

    private static final Logger logger = LoggerFactory.getLogger(TimestampConsumer.class);
    private static String bootstrapServer = null;
    private static String topic = null;
    private static String saslUsername = null;
    private static String saslPassword = null;
    private static String krb5Conf = null;
    private static String jaasConf = null;
    private static String startTime = null;
    private static String groupId = "test-consumer-group";
    private static String interval = null;
    private static boolean isDebug = false;
    private static boolean isPrintTimestamp = false;
    private static long startTimestamp;
    private static long endTimestamp;
    private static long beginningTime = -1;
    private static long intervalTime;
    private static String offsetFile;
    private static String offsetDate;
    private static ConsumerType type;
    private static final FastDateFormat DATE_FORMAT = FastDateFormat.getInstance("yyyy-MM-dd HH:mm:ss");
    private static boolean isJarRun = false;


    private static void init() {
        bootstrapServer = "10.243.141.133:9092";
        topic = "test";
        startTime = "earliest";
        isPrintTimestamp = true;

    }

    public static void main(String[] args) throws Exception {
        if (isJarRun) {
            parseArgs(args);
        } else {
            init();
        }
        initValue();
        Properties properties = new Properties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer);
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true);

        if (StringUtils.isNoneEmpty(krb5Conf, jaasConf)) {
            logger.info("start build Sasl env");
            KafkaAuth.buildKafkaSaslEnv(properties, krb5Conf, jaasConf);
        }
        if (StringUtils.isNoneEmpty(saslUsername, saslPassword)) {
            logger.info("start build kerberos env");
            KafkaAuth.buildKafkaSaslEnv(properties, saslUsername, saslPassword);
        }

        KafkaConsumer<String, String> consumer = buildKafkaConsumer(properties);

        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(1000);
            if (records.isEmpty() && type != ConsumerType.LATEST) {
                System.exit(0);
            }
            for (ConsumerRecord<String, String> record : records) {
                if (beginningTime < 0) {
                    if (ConsumerType.TIMESTAMP == type) {
                        beginningTime = startTimestamp;
                    } else {
                        beginningTime = record.timestamp();
                    }
                    if (intervalTime > 0) {
                        endTimestamp = beginningTime + intervalTime;
                    } else {
                        SimpleDateFormat nextDateFormat = new SimpleDateFormat("yyyy-MM-dd 00:00:00");
                        Date nextDate = DateUtils.addDays(new Date(beginningTime), 1);
                        String format = nextDateFormat.format(nextDate);
                        endTimestamp = nextDateFormat.parse(format).getTime();
                    }
                    logger.info("kafka consumer info: type:{},topic:{},start-time:{},interval:{},end-time:{}",
                            type,
                            topic,
                            DATE_FORMAT.format(beginningTime),
                            interval,
                            DATE_FORMAT.format(endTimestamp)
                    );
                }
                if (record.timestamp() > endTimestamp) {
                    System.exit(0);
                }
                if (isPrintTimestamp) {
                    System.out.println(String.format("[%s, %s] %s", new Timestamp(record.timestamp()).toString(), record.key(), record.value()));
                } else {
                    System.out.println(record.value());
                }
                if (isDebug) {
                    Thread.sleep(1000);
                }
            }
        }
    }

    private static KafkaConsumer buildKafkaConsumer(Properties properties) {
        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(properties);
        List<PartitionInfo> partitionSet = consumer.partitionsFor(topic);
        if (isDebug) {
            logger.info("Partition Info {}", partitionSet.toString());
        }

        switch (type) {
            case TIMESTAMP:
                Map<TopicPartition, Long> map = new HashMap<>(16);
                for (PartitionInfo partitionInfo : partitionSet) {
                    TopicPartition topicPartition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
                    map.put(topicPartition, startTimestamp);
                }
                Map<TopicPartition, OffsetAndTimestamp> partitionOffsetAndTimestampMap = consumer.offsetsForTimes(map);

                if (partitionOffsetAndTimestampMap.isEmpty()) {
                    System.err.println("get offset from timestamp error,data not found");
                    throw new RuntimeException("data not found");
                }
                if (isDebug) {
                    logger.info("Partition Offset And Timestamp {}", partitionOffsetAndTimestampMap.toString());
                }
                for (Map.Entry<TopicPartition, OffsetAndTimestamp> entry : partitionOffsetAndTimestampMap.entrySet()) {
                    consumer.assign(Collections.singletonList(entry.getKey()));
                    consumer.seek(entry.getKey(), entry.getValue().offset());
                }
                break;
            case LATEST:
                for (PartitionInfo partitionInfo : partitionSet) {
                    TopicPartition topicPartition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
                    consumer.assign(Collections.singletonList(topicPartition));
                    consumer.seekToEnd(Collections.singletonList(topicPartition));
                }
                break;
            case EARLIEST:
                for (PartitionInfo partitionInfo : partitionSet) {
                    TopicPartition topicPartition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
                    consumer.assign(Collections.singletonList(topicPartition));
                    consumer.seekToBeginning(Collections.singletonList(topicPartition));
                }
                break;
            case OFFSET:
                if (isDebug) {
                    logger.info("Partition Offset {}", parseOffset().toString());
                }
                for (Map.Entry<TopicPartition, OffsetAndTimestamp> entry : parseOffset().entrySet()) {
                    consumer.assign(Collections.singletonList(entry.getKey()));
                    consumer.seek(entry.getKey(), entry.getValue().offset());
                }
                break;
            default:
                for (PartitionInfo partitionInfo : partitionSet) {
                    TopicPartition topicPartition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
                    consumer.assign(Collections.singletonList(topicPartition));
                    consumer.seekToEnd(Collections.singletonList(topicPartition));
                }
                break;
        }
        return consumer;
    }

    private static HashMap<TopicPartition, OffsetAndTimestamp> parseOffset() {
        if (null == offsetDate) {
            Date today = new Date(System.currentTimeMillis() - 1000 * 60 * 60 * 24);
            SimpleDateFormat simpleDateFormat = new SimpleDateFormat("yyyyMMdd");
            offsetDate = simpleDateFormat.format(today);
        }
        File f = new File(offsetFile);
        if (!f.exists()) {
            System.out.println(String.format("%s not exist!", offsetFile));
            System.exit(0);
        }
        HashMap<TopicPartition, OffsetAndTimestamp> topicPartitionMap = new HashMap<>(1);
        BufferedReader reader = null;
        try {
            reader = new BufferedReader(new FileReader(f));
            String line;
            while ((line = reader.readLine()) != null) {
                if (line.startsWith(offsetDate)) {
                    String[] split = line.split(":");
                    String t = split[1];
                    if (topic.equals(t)) {
                        int partition = Integer.parseInt(split[2]);
                        long offset = Long.parseLong(split[3]);
                        TopicPartition topicPartition = new TopicPartition(topic, partition);
                        OffsetAndTimestamp offsetAndTimestamp = new OffsetAndTimestamp(offset, -1);
                        topicPartitionMap.put(topicPartition, offsetAndTimestamp);
                    }
                }
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
        }

        return topicPartitionMap;
    }


    /**
     * 解析命令参数
     *
     * @param args 参数
     */
    private static void parseArgs(String[] args) throws Exception {
        CommandLineParser parser = new DefaultParser();
        Options options = new Options();
        options.addOption("topic", "topic", true, "topic");
        options.addOption("t", "time", true, "start-time");
        options.addOption("bs", "bootstrap-server", true, "bootstrap-server");
        options.addOption("kf", "krb5-conf", true, "krb5-conf file path");
        options.addOption("jf", "jaas-conf", true, "jaas-conf file path");
        options.addOption("su", "sasl-username", true, "sasl username");
        options.addOption("sp", "sasl-password", true, "sasl password");
        options.addOption("gi", "group-id", true, "group id");
        options.addOption("h", "help", false, "help info");
        options.addOption("i", "interval", true, "interval");
        options.addOption("date", "date", true, "date");
        options.addOption("d", "debug", false, "debug");
        options.addOption("pt", "print-timestamp", false, "print timestamp");

        CommandLine commandLine = parser.parse(options, args, true);


        if (commandLine.hasOption("help") || args.length == 0) {
            HelpFormatter formatter = new HelpFormatter();
            formatter.printHelp("java -jar", options);
            System.exit(0);
        }

        if (commandLine.hasOption(CliConstant.BOOTSTRAP_SERVERS)) {
            bootstrapServer = commandLine.getOptionValue(CliConstant.BOOTSTRAP_SERVERS);
        }
        if (commandLine.hasOption(CliConstant.TOPIC)) {
            topic = commandLine.getOptionValue(CliConstant.TOPIC);
        }
        if (commandLine.hasOption(CliConstant.TIME)) {
            startTime = commandLine.getOptionValue(CliConstant.TIME).trim();
        }
        if (commandLine.hasOption(CliConstant.DATE)) {
            offsetDate = commandLine.getOptionValue(CliConstant.DATE).trim();
        }
        if (commandLine.hasOption(CliConstant.INTERVAL)) {
            interval = commandLine.getOptionValue(CliConstant.INTERVAL).trim();
        }
        if (commandLine.hasOption(CliConstant.IS_PRINT_TIMESTAMP)) {
            isPrintTimestamp = true;
        }
        if (commandLine.hasOption(CliConstant.SASL_USERNAME)) {
            saslUsername = commandLine.getOptionValue(CliConstant.SASL_USERNAME);
        }
        if (commandLine.hasOption(CliConstant.SASL_PASSWORD)) {
            saslPassword = commandLine.getOptionValue(CliConstant.SASL_PASSWORD);
        }
        if (commandLine.hasOption(CliConstant.KRB5_CONFIG)) {
            krb5Conf = commandLine.getOptionValue(CliConstant.KRB5_CONFIG);
        }
        if (commandLine.hasOption(CliConstant.JAAS_CONFIG)) {
            jaasConf = commandLine.getOptionValue(CliConstant.JAAS_CONFIG);
        }
        if (commandLine.hasOption(CliConstant.GROUP_ID)) {
            groupId = commandLine.getOptionValue(CliConstant.GROUP_ID).trim();
        }
        if (commandLine.hasOption(CliConstant.IS_DEBUG)) {
            isDebug = true;
        }

    }

    /**
     * 初始化value
     */
    private static void initValue() throws Exception {
        if (StringUtils.isAnyEmpty(bootstrapServer, topic, startTime)) {
            System.err.println("argument bootstrap-server topic time can not be null");
            throw new MissingArgumentException("miss argument");
        }

        if (CliConstant.EARLIEST.equals(startTime)) {
            type = ConsumerType.EARLIEST;
        } else if (CliConstant.LATEST.equals(startTime)) {
            type = ConsumerType.LATEST;
        } else {
            try {
                long l = Long.parseLong(startTime);
                if (l > 0) {
                    SimpleDateFormat timeFormat = new SimpleDateFormat("yyyyMMddHHmmss");
                    Date startDate = timeFormat.parse(startTime);
                    startTimestamp = startDate.getTime();
                    type = ConsumerType.TIMESTAMP;
                } else {
                    if (CliConstant.LATEST_NUM == l) {
                        type = ConsumerType.LATEST;
                    } else if (CliConstant.EARLIEST_NUM == l) {
                        type = ConsumerType.EARLIEST;
                    } else {
                        throw new IllegalArgumentException("invalid input time :" + startTime);
                    }
                }
            } catch (Exception e) {
                offsetFile = startTime;
                type = ConsumerType.OFFSET;
            }
        }
        if (null != interval) {
            if (interval.endsWith(CliConstant.MINUTE)) {
                intervalTime = Long.parseLong(interval.replace("m", "")) * 60 * 1000;
            } else {
                intervalTime = Long.parseLong(interval.replace("h", "")) * 60 * 60 * 1000;
            }
        }

        if (isDebug) {
            logAll();
        }
    }

    private static void logAll() {
        System.out.println(String.format(
                "=======================================\n" +
                        "Kafka consumer config:\n" +
                        "consumer.type:%s\n" +
                        "bootstrap.servers:%s\n" +
                        "topic:%s\n" +
                        "group-id：%s\n" +
                        "start.time:%s\n" +
                        "interval:%s\n" +
                        "sasl-username:%s\n" +
                        "sasl-password:%s\n" +
                        "krb5-conf:%s\n" +
                        "jaas-conf:%s\n" +
                        "=======================================\n"
                , type, bootstrapServer, topic, groupId, startTime, interval, saslUsername, saslPassword, krb5Conf, jaasConf));
    }

    enum ConsumerType {
        /**
         * 消费类型
         */
        TIMESTAMP,
        LATEST,
        EARLIEST,
        OFFSET
    }

}
